from confluent_kafka import Producer, Consumer
from confluent_kafka.admin import AdminClient, NewTopic


class Kafka:

    DefaultConfigure = {}

    def init_from_params(self, **kwargs):
        if not kwargs.get('bootstrap.servers'):
            host = kwargs.pop('host')
            port = kwargs.pop('port')
            kwargs['bootstrap.servers'] = f'{host}:{port}'
        return {**self.DefaultConfigure, **kwargs}


class KafkaProducer(Kafka):

    DefaultConfigure = {'linger.ms': 500, 'batch.size': 1000, 'acks': 1}

    def __init__(self, **kwargs):
        self.producer = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        kwargs = self.init_from_params(**kwargs)
        self.producer = Producer(**kwargs)

    def send(self, topic, value, *args, **kwargs):
        try:
            self.producer.produce(topic, value, *args, **kwargs)
        except BufferError:
            self.producer.flush()


class KafkaConsumer(Kafka):

    DefaultConfigure = {'default.topic.config': {'auto.offset.reset': 'latest'}}

    def __init__(self, **kwargs):
        self.consumer = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        kwargs = self.init_from_params(**kwargs)
        self.consumer = Consumer(kwargs)

    def poll(self, timeout=1):
        return self.consumer.poll(timeout)

    def consume(self, num_messages=1, timeout=-1):
        return self.consumer.consume(num_messages=num_messages, timeout=timeout)


class KafkaAdmin(Kafka):
    def __init__(self, **kwargs):
        self.admin = None
        if kwargs:
            self.init(**kwargs)

    def init(self, **kwargs):
        kwargs = self.init_from_params(**kwargs)
        self.admin = AdminClient(kwargs)

    def list_topics(self):
        return self.admin.list_topics()

    def create_topics(self, topics):
        if isinstance(topics, str):
            topics = [topics]
        return self.admin.create_topics([NewTopic(item) for item in topics])

    def delete_topics(self, topics):
        if isinstance(topics, str):
            topics = [topics]
        return self.admin.delete_topics(topics)


kafka_producer = KafkaProducer()
kafka_consumer = KafkaConsumer()
kafka_admin = KafkaAdmin()


if __name__ == '__main__':
    kafka_admin.init(**{'bootstrap.servers': '192.168.101.29:9092'})
    kafka_producer.init(**{'bootstrap.servers': '192.168.101.29:9092'})
    kafka_consumer.init(**{'bootstrap.servers': '192.168.101.29:9092', 'group.id': 'b',
                           'default.topic.config': {'auto.offset.reset': 'smallest'}})
    kafka_consumer.consumer.subscribe(['a10'])


    def test_admin():
        result = kafka_admin.admin.list_topics()
        print(result.topics)

    def test_sender():
        kafka_producer.send('a10', '11')

    def test_consumer():
        message = kafka_consumer.poll(1)
        messages = kafka_consumer.consume()


    test_admin()
    test_sender()
    test_consumer()
